/*
 * Copyright (c) <2010>, <lishun>
 * All rights reserved.
 *
 * Redistribution and use in source and binary forms, with or without
 * modification, are permitted provided that the following conditions
 * are met:
 * 1. Redistributions of source code must retain the above copyright
 *    notice, this list of conditions and the following disclaimer.
 * 2. Redistributions in binary form must reproduce the above copyright
 *    notice, this list of conditions and the following disclaimer in the
 *    documentation and/or other materials provided with the distribution.
 * 4. Neither the name of the University nor the names of its contributors
 *    may be used to endorse or promote products derived from this software
 *    without specific prior written permission.
 *
 * THIS SOFTWARE IS PROVIDED BY THE REGENTS AND CONTRIBUTORS ``AS IS'' AND
 * ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
 * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
 * ARE DISCLAIMED.  IN NO EVENT SHALL THE REGENTS OR CONTRIBUTORS BE LIABLE
 * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
 * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS
 * OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION)
 * HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT
 * LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY
 * OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF
 * SUCH DAMAGE.
 *
 * $Id$
*/

#include <sys/limits.h>
#include <sys/queue.h>
#include <sys/types.h>

#include <stddef.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>

struct ahoc_out_t {
	SLIST_ENTRY(ahoc_out_t) link;
	void *tag;
	int len;
};
SLIST_HEAD(ahoc_outhead, ahoc_out_t);

struct ahoc_node_t {
	TAILQ_ENTRY(ahoc_node_t) link;
	struct ahoc_node_t **go;
	uint16_t gochar[2];
	struct ahoc_node_t *goarr[2];
	struct ahoc_node_t *fail;
	struct ahoc_outhead outq;
};
TAILQ_HEAD(ahoc_nodehead, ahoc_node_t);

typedef int (*ahoc_outfunc)(int index, void *tag);
struct ahoc_automaton {
	struct ahoc_node_t root;
	ahoc_outfunc outfunc;
	int ignore_case;
	int memuse;
};

#define AHOC_GOTO(var, c)	(var->go ? var->go[c] : \
				 var->gochar[0] == c ? var->goarr[0] :\
				 var->gochar[1] == c ? var->goarr[1] :\
				 NULL)

char ahoc_lowcase[UCHAR_MAX];
static int ahoc_case_inited = 0;

void
ahoc_caseinit()
{
	int i;
	for (i = 0; i < UCHAR_MAX; i++) {
		if (i >= 'A' && i <= 'Z')
			ahoc_lowcase[i] = i | 0x20;
		else
			ahoc_lowcase[i] = i;
	}
	ahoc_case_inited = 1;

}

struct ahoc_automaton *
ahoc_new(ahoc_outfunc func, int ignore_case)
{
	struct ahoc_automaton *ac;

	if (ahoc_case_inited == 0)
		ahoc_caseinit();

	ac = malloc(sizeof(struct ahoc_automaton));
	ac->memuse = sizeof(struct ahoc_automaton);
	ac->root.go = calloc(UCHAR_MAX, sizeof(void *));
	ac->memuse += UCHAR_MAX * sizeof(void *);

	ac->root.gochar[0] = USHRT_MAX;
	ac->root.gochar[1] = USHRT_MAX;

	ac->root.fail = &ac->root;
	SLIST_INIT(&ac->root.outq);
	ac->outfunc = func;
	ac->ignore_case = ignore_case;
	return ac;
}

void
ahoc_add(struct ahoc_automaton *ac, uint8_t *buf, int len, void *tag)
{
	struct ahoc_node_t *state;
	struct ahoc_node_t *node;
	struct ahoc_out_t *out;
	int i;
	uint8_t c;

	state = &ac->root;
	for (i = 0; i < len; i++) {
		c = buf[i];
		if (ac->ignore_case)
			c = ahoc_lowcase[c];
		if (AHOC_GOTO(state, c) != NULL) {
			state = AHOC_GOTO(state, c);
			continue;
		}
		node = malloc(sizeof(struct ahoc_node_t));
		ac->memuse += sizeof(struct ahoc_node_t);
		node->go = NULL;
		node->gochar[0] = USHRT_MAX;
		node->gochar[1] = USHRT_MAX;
		node->fail = NULL;
		SLIST_INIT(&node->outq);

		if (state->go != NULL)
			state->go[c] = node;
		else if (state->gochar[0] != USHRT_MAX &&
		    state->gochar[1] != USHRT_MAX) {
			state->go = calloc(UCHAR_MAX, sizeof(void *));
			ac->memuse += UCHAR_MAX * sizeof(void *);
			state->go[c] = node;
			state->go[state->gochar[0]] = state->goarr[0];
			state->go[state->gochar[1]] = state->goarr[1];
		} else if (state->gochar[0] == USHRT_MAX) {
			state->gochar[0] = c;
			state->goarr[0] = node;
		} else {
			state->gochar[1] = c;
			state->goarr[1] = node;
		}

		state = node;
	}
	out = malloc(sizeof(struct ahoc_out_t));
	ac->memuse += sizeof(struct ahoc_out_t);
	out->tag = tag;
	out->len = len;
	SLIST_INSERT_HEAD(&state->outq, out, link);
}

void
ahoc_prepare(struct ahoc_automaton *ac)
{
	int i;
	struct ahoc_nodehead queue;
	struct ahoc_node_t *r, *s, *t;
	struct ahoc_out_t *o, *p;

	for (i = 0; i < UCHAR_MAX; i++) {
		if (ac->root.go[i] == NULL)
			ac->root.go[i] = &ac->root;
	}

	TAILQ_INIT(&queue);
	for (i = 0; i < UCHAR_MAX; i++) {
		s = ac->root.go[i];
		if (s != &ac->root) {
			TAILQ_INSERT_TAIL(&queue, s, link);
			s->fail = &ac->root;
		}
	}
	while (!TAILQ_EMPTY(&queue)) {
		r = TAILQ_FIRST(&queue);
		TAILQ_REMOVE(&queue, r, link);
		for (i = 0; i < UCHAR_MAX; i++) {
			s = AHOC_GOTO(r, i);
			if (s == NULL)
				continue;
			TAILQ_INSERT_TAIL(&queue, s, link);
			t = r->fail;
			while (AHOC_GOTO(t, i) == NULL)
				t = t->fail;
			s->fail = AHOC_GOTO(t, i);
			SLIST_FOREACH(o, &s->fail->outq, link) {
				p = malloc(sizeof(struct ahoc_out_t));
				ac->memuse += sizeof(struct ahoc_out_t);
				p->tag = o->tag;
				p->len = o->len;
				SLIST_INSERT_HEAD(&s->outq, p, link);
			}
		}
	}
}

void
ahoc_search(struct ahoc_automaton *ac, const uint8_t *buf, uint32_t len)
{
	struct ahoc_node_t *state;
	struct ahoc_out_t *out;
	int i;
	uint8_t c;
	struct ahoc_node_t *oldmatch = NULL;

	state = &ac->root;
	for (i = 0; i < len; i++) {
		if (ac->ignore_case)
			c = ahoc_lowcase[buf[i]];
		else
			c = buf[i];
		while (AHOC_GOTO(state, c) == NULL)
			state = state->fail;
		state = AHOC_GOTO(state, c);
		if (SLIST_EMPTY(&state->outq))
			continue;
		if (oldmatch == state)
			continue;
		oldmatch = state;
		SLIST_FOREACH(out, &state->outq, link) {
			int ret = 0;
			if (ac->outfunc != NULL)
				ret = ac->outfunc(i - out->len + 1, out->tag);
			if (ret != 0)
				return;
		}
	}
}

void
ahoc_free(struct ahoc_automaton *ac)
{
	int i;
	struct ahoc_nodehead queue;
	struct ahoc_node_t *node;

	TAILQ_INIT(&queue);
	for (i = 0; i < UCHAR_MAX; i++) {
		node = ac->root.go[i];
		if (node != &ac->root)
			TAILQ_INSERT_TAIL(&queue, node, link);
	}
	while (!TAILQ_EMPTY(&queue)) {
		node = TAILQ_FIRST(&queue);
		TAILQ_REMOVE(&queue, node, link);
		for (i = 0; i < UCHAR_MAX; i++) {
			struct ahoc_node_t *tnode = AHOC_GOTO(node, i);
			if (tnode != 0)
				TAILQ_INSERT_TAIL(&queue, tnode, link);
		}

		while (!SLIST_EMPTY(&node->outq)) {
			struct ahoc_out_t *out = SLIST_FIRST(&node->outq);
			SLIST_REMOVE_HEAD(&node->outq, link);
			free(out);
		}
		if (node->go)
			free(node->go);
		free(node);
	}
	free(ac);
}

#if 1
int
ahoc_printtag(int index, void *tag)
{
	printf("%d %c\n", index, (int)tag);
	return 0;
}

int main()
{
	struct ahoc_automaton *ac;
	FILE *fp;
	char txt[1024];
	int n, i;

	fp = fopen("./ahoc.c", "r");
	n = fread(txt, 1, 1024, fp);
	fclose(fp);

	ac = ahoc_new(ahoc_printtag, 1);

	ahoc_add(ac, "HE", 2, (void *)'h');
	ahoc_add(ac, "SHE", 3, (void *)'s');
	ahoc_add(ac, "HIS", 3, (void *)'i');
	ahoc_add(ac, "HERS", 4, (void *)'r');
	ahoc_add(ac, "struct", 6, (void *)'t');

	ahoc_prepare(ac);

	for (i=0; i<100000; i++) {
		ahoc_search(ac, txt, n);
	}

	ahoc_free(ac);
	return 0;
}
#endif
